from langchain_community.llms.tongyi import Tongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# RAG增强

from langchain_community.vectorstores import FAISS
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.embeddings.dashscope import DashScopeEmbeddings


llm = Tongyi()
# prompt = ChatPromptTemplate.from_template("给我讲一个关于{topic}的故事,返回字数不超过50个字")

# prompt_value = prompt.invoke({"topic":"兔子"})
# print(prompt_value)
# print(prompt_value.to_messages())
# print(prompt_value.to_string())

# out_parser = StrOutputParser()
# chain = prompt | llm | out_parser
# res = chain.invoke({"topic":"兔子"})
# print(res)

# RAG增强

db = FAISS.from_texts(
    # 向量数据库中的文本数据
    ["张三在北京工作,一个月工资3万人民币"],
    # 向量嵌入器
    embedding = DashScopeEmbeddings(),
)

# 将向量存储转换为检索器
reteiever = db.as_retriever()

template = """根据下面内容回答问题
{context}
问题: {question}
"""

prompt = ChatPromptTemplate.from_template(template)

# RunnablePassthrough 被用来确保问题直接传递到下一个组件,不做任何处理,确保问题被注解传递 RunnableParallel 的输出中,然后可以被 prompt 使用
# RunnableParallel 允许开发者并行运行多个可运行对象,并将结果合并,并行处理上下文核问题

chain = (RunnableParallel({
    # 检索器
    "context":reteiever,
    # 模板
    "question":RunnablePassthrough()})
    # 将上下文和问题作为输入传递给提示模板
    | prompt
    # 将模板输出提示词传递给LLM
    | llm
    # 将LLM输出解析成字符串后输出
    | StrOutputParser()
)

res = chain.invoke(input="张三在北京工作,一个月工资多少?")
print(res)
